(self)
| 82 | return (2,) * int(math.log2(target_codebook_size)) |
| 83 | |
| 84 | def setup(self): |
| 85 | self.proj_down = nn.Dense(len(self.bins_per_dim)) |
| 86 | self.proj_up = nn.Dense(self.input_dim) |
| 87 | |
| 88 | def __call__(self, inputs: jnp.ndarray) -> tuple[jnp.ndarray, jnp.ndarray]: |
| 89 | tokens, z = self.encode(inputs) |
nothing calls this directly
no outgoing calls
no test coverage detected