(self, inputs: jnp.ndarray)
| 91 | return tokens, output |
| 92 | |
| 93 | def encode(self, inputs: jnp.ndarray) -> tuple[jnp.ndarray, jnp.ndarray]: |
| 94 | bases = jnp.array(self.bins_per_dim) |
| 95 | |
| 96 | x = self.proj_down(inputs) |
| 97 | z = jnp.tanh(x) |
| 98 | |
| 99 | # Quantize |
| 100 | digits = jnp.round((z + 1) * (bases - 1) / 2).astype(jnp.int32) |
| 101 | tokens = self.undigitize(digits) |
| 102 | |
| 103 | return tokens, z |
| 104 | |
| 105 | def decode(self, tokens: jnp.ndarray, z_grad: jax.Array | None = None) -> jnp.ndarray: |
| 106 | bases = jnp.array(self.bins_per_dim) |
no test coverage detected