Updates the code to always create variables and computations of the same dtype as the its inputs. Previously, if float64 was enabled, some of the results would be (potentially incorrectly) promoted to higher precision.

kfac_jax/_src/curvature_blocks.py

@@ -201,10 +201,16 @@ def inputs_shapes(self) -> Tuple[Shape, ...]:
   @property
   def parameters_shapes(self) -> Tuple[Shape, ...]:
     """The shapes of the parameter variables of the block's tag equation."""
-
     return tuple(jax.tree_util.tree_map(
         lambda x: tuple(x.aval.shape), self.parameter_variables))
 
+  @property
+  def dtype(self) -> chex.ArrayDType:
+    dtypes = set(p.aval.dtype for p in self.parameter_variables)  # pytype: disable=attribute-error
+    if len(dtypes) > 1:
+      raise ValueError("Not all parameters are the same dtype.")
+    return dtypes.pop()
+
   @property
   def parameters_canonical_order(self) -> Tuple[int, ...]:
     """The canonical order of the parameter variables."""
@@ -666,8 +672,8 @@ def _init(
 
     return Diagonal.State(
         cache=None,
-        diagonal_factors=tuple(utils.WeightedMovingAverage.zero(s)
-                               for s in self.parameters_shapes),
+        diagonal_factors=tuple(utils.WeightedMovingAverage.zero(
+            shape, self.dtype) for shape in self.parameters_shapes),
     )
 
   def _multiply_matpower_unscaled(
@@ -854,18 +860,20 @@ def _init(
     cache = {}
 
     if len(exact_powers_to_cache) > self._eigen_decomposition_threshold:
-      cache["eigenvalues"] = jnp.zeros([self.dim])
-      cache["eigen_vectors"] = jnp.zeros([self.dim, self.dim])
+      cache["eigenvalues"] = jnp.zeros([self.dim], self.dtype)
+      cache["eigen_vectors"] = jnp.zeros([self.dim, self.dim], self.dtype)
+
     elif cache_eigenvalues:
-      cache["eigenvalues"] = jnp.zeros([self.dim])
+      cache["eigenvalues"] = jnp.zeros([self.dim], self.dtype)
 
     if len(exact_powers_to_cache) <= self._eigen_decomposition_threshold:
       for power in exact_powers_to_cache:
-        cache[str(power)] = jnp.zeros([self.dim, self.dim])
+        cache[str(power)] = jnp.zeros([self.dim, self.dim], self.dtype)
 
     return Full.State(
         cache=cache,
-        matrix=utils.WeightedMovingAverage.zero((self.dim, self.dim)),
+        matrix=utils.WeightedMovingAverage.zero(
+            [self.dim, self.dim], self.dtype),
     )
 
   def _multiply_matpower_unscaled(
@@ -978,8 +986,8 @@ def _update_cache(
     else:
 
       if eigenvalues:
-        state.cache["eigenvalues"] = (
-            scale * utils.safe_psd_eigh(state.matrix.value)[0])
+        state.cache["eigenvalues"] = scale * utils.safe_psd_eigh(
+            state.matrix.value)[0]
 
       for power in exact_powers:
 
@@ -1078,26 +1086,29 @@ def _init(
     cache = {}
 
     if cache_eigenvalues or exact_powers_to_cache:
-      cache["inputs_factor_eigenvalues"] = jnp.zeros([d_in])
-      cache["outputs_factor_eigenvalues"] = jnp.zeros([d_out])
+      cache["inputs_factor_eigenvalues"] = jnp.zeros([d_in], self.dtype)
+      cache["outputs_factor_eigenvalues"] = jnp.zeros([d_out], self.dtype)
 
     if exact_powers_to_cache:
-      cache["inputs_factor_eigen_vectors"] = jnp.zeros([d_in, d_in])
-      cache["outputs_factor_eigen_vectors"] = jnp.zeros([d_out, d_out])
+      cache["inputs_factor_eigen_vectors"] = jnp.zeros([d_in, d_in], self.dtype)
+      cache["outputs_factor_eigen_vectors"] = jnp.zeros(
+          [d_out, d_out], self.dtype)
 
     for power in approx_powers_to_cache:
       if power != -1:
         raise NotImplementedError(f"Approximations for power {power} is not "
                                   f"yet implemented.")
       cache[str(power)] = dict(
-          inputs_factor=jnp.zeros([d_in, d_in]),
-          outputs_factor=jnp.zeros([d_out, d_out]),
+          inputs_factor=jnp.zeros([d_in, d_in], self.dtype),
+          outputs_factor=jnp.zeros([d_out, d_out], self.dtype),
       )
 
     return TwoKroneckerFactored.State(
         cache=cache,
-        inputs_factor=utils.WeightedMovingAverage.zero((d_in, d_in)),
-        outputs_factor=utils.WeightedMovingAverage.zero((d_out, d_out)),
+        inputs_factor=utils.WeightedMovingAverage.zero(
+            [d_in, d_in], self.dtype),
+        outputs_factor=utils.WeightedMovingAverage.zero(
+            [d_out, d_out], self.dtype),
     )
 
   def _multiply_matpower_unscaled(
@@ -1226,7 +1237,6 @@ def _update_cache(
     factor_scale = jnp.power(scale, 0.5)
 
     if eigenvalues or exact_powers:
-
       s_i, q_i = utils.safe_psd_eigh(state.inputs_factor.value)
       s_o, q_o = utils.safe_psd_eigh(state.outputs_factor.value)
 

kfac_jax/_src/utils/accumulators.py

@@ -56,10 +56,14 @@ def sync(self, pmap_axis_name: Optional[str]) -> None:
     self.raw_value = parallel.pmean_if_pmap(self.raw_value, pmap_axis_name)
 
   @classmethod
-  def zero(cls, shape: chex.Shape) -> "WeightedMovingAverage":
+  def zero(
+      cls,
+      shape: chex.Shape,
+      dtype: Optional[chex.ArrayDType] = None,
+  ) -> "WeightedMovingAverage":
     """Initializes a `WeightedMovingAverage` with a single array of zeros."""
     return WeightedMovingAverage(
-        weight=jnp.zeros([]), raw_value=jnp.zeros(shape))
+        weight=jnp.zeros([], dtype), raw_value=jnp.zeros(shape, dtype))
 
   @classmethod
   def zeros_like(cls, value: PyTree) -> "WeightedMovingAverage":

kfac_jax/_src/utils/math.py

@@ -324,7 +324,7 @@ def psd_inv_cholesky(matrix: chex.Array, damping: chex.Array) -> chex.Array:
   if matrix.shape[:1] != matrix.shape[1:]:
     raise ValueError(f"Expected square matrix, but got shape {matrix.shape}.")
 
-  identity = jnp.eye(matrix.shape[0])
+  identity = jnp.eye(matrix.shape[0], dtype=matrix.dtype)
 
   return linalg.solve(matrix + damping * identity, identity, assume_a="pos")
 
@@ -377,6 +377,7 @@ def pi_adjusted_kronecker_inverse(
 
   # kron(arrays) = c * kron(us)
   c = jnp.exp(jnp.sum(jnp.log(jnp.stack(norms)) - jnp.log(jnp.stack(dims))))
+  damping = damping.astype(c.dtype)
 
   def regular_inverse() -> Tuple[chex.Array, ...]:
 
@@ -417,7 +418,7 @@ def zero_inverse() -> Tuple[chex.Array, ...]:
     for a in us:
 
       if a.ndim == 2:
-        inv = jnp.eye(a.shape[0])
+        inv = jnp.eye(a.shape[0], dtype=a.dtype)
 
       else:
         inv = jnp.ones_like(a)
@@ -643,7 +644,8 @@ def safe_psd_eigh(
   # of cuda and cudablas they can cause a runtime error.
   s, q = lax.cond(
       jnp.any(jnp.isnan(x)),
-      lambda _: (jnp.full([d], jnp.nan), jnp.full([d, d], jnp.nan)),
+      lambda _: (jnp.full([d], jnp.nan, dtype=x.dtype),  # pylint: disable=g-long-lambda
+                 jnp.full([d, d], jnp.nan, dtype=x.dtype)),
       functools.partial(_eigh, force_on_host=force_on_host),
       x,
   )