Pytorch-эквивалент `register_buffer` в Flax/jax ⇐ Python

[Anonymous]

Я ищу способ написать эквивалент следующего модуля Pytorch во Flax, но не нашел способа это сделать. Важно то, что константа должна быть загружаемой и сохраняемой при контрольной точке.

Код: Выделить всё

class SillyModule(nn.Module):
def __init__(self, ):
super().__init__()
self.register_buffer('constant', torch.randn(1, 128))

def forward(self, x):
return torch.matmul(x, self.B)

Кто-нибудь знает, как это сделать? Что является эквивалентом Register_buffer во льне?

Подробнее здесь: https://stackoverflow.com/questions/732 ... n-flax-jax