(self, x: mx.array)
| 345 | return new_weights |
| 346 | |
| 347 | def encode(self, x: mx.array): |
| 348 | z = self.reg(self.encoder(x)) |
| 349 | z = self.scale_factor * (z - self.shift_factor) |
| 350 | return z |
| 351 | |
| 352 | def decode(self, z: mx.array): |
| 353 | z = z / self.scale_factor + self.shift_factor |