from json import dumps
from torch import autograd, nn, tensor
from arti import mechanisms as api

pair = api.TensorType(("C",), (2,), dtype="float32")
signal = api.InputBinding("signal", pair)
carry = api.InputBinding("carry", pair)
bank = api.BankBinding("weight", "tutorial/carry-weight@1", "weight", api.TensorType((), (), dtype="float32"))
transition = api.add(carry, api.scale(signal, bank))
summary = api.reduce_tensor(transition, axis="C", mode="sum")
program = api.FormulaProgram.build(outputs=(summary, transition))
parameter = nn.Parameter(tensor(0.75))
provided = dict(signal=tensor([4.0, -2.0]), carry=tensor([1.0, 1.0]))
emitted = api.FormulaFabricV2(program)(inputs=provided, banks=dict(weight=bank.bind(parameter)))
readout, returned_carry = emitted.values
gradient = autograd.grad(readout, parameter)[0]
record = dict(readout=readout.item(), next_carry=returned_carry.detach().tolist(), weight_vjp=gradient.item())
print(dumps(record))