File size: 901 Bytes
f118c1e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 |
"""Test submission for ARC using identity featurizer."""
import torch
import torch.nn as nn
import pyvene as pv
from CausalAbstraction.neural.featurizers import Featurizer
class IdentityFeaturizerModule(torch.nn.Module):
def __init__(self):
super().__init__()
self.dummy_param = torch.nn.Parameter(torch.zeros(1), requires_grad=True)
def forward(self, x):
# return x, None
return x + 0 * self.dummy_param.sum(), None
class IdentityInverseFeaturizerModule(torch.nn.Module):
def __init__(self):
super().__init__()
def forward(self, f, error):
return f
class IdentityFeaturizer(Featurizer):
def __init__(self, id="identity"):
featurizer = IdentityFeaturizerModule()
inverse_featurizer = IdentityInverseFeaturizerModule()
super().__init__(featurizer, inverse_featurizer, id=id) |