How can I make my function bfloat16 compatible?

[pvasired]

I have the following function which I am trying to speed up taking gradients of using the bfloat16 data type.

@jax.jit
def fisher_loss_array(probs_vec, transform_mat, jac_full, trials):
    """
    Compute the Fisher loss across the entire array.

    Parameters:
    probs_vec (jnp.DeviceArray): The flattened array of probabilities across all meaningful
                            (cell, pattern) combinations
    transform_mat (jnp.DeviceArray): The transformation matrix to convert the trials array
                                to the transformed trials array for multiple cells on
                                the same pattern
    jac_full (jnp.DeviceArray): The full precomputed Jacobian matrix
    trials (jnp.DeviceArray): The input trials vector to be optimized

    Returns:
    loss (float): The whole array Fisher information loss
    """
    p_model = jnp.clip(probs_vec, a_min=1e-5, a_max=1-1e-5) # need to clip these to prevent
                                                            # overflow errors
    t = jnp.dot(transform_mat, trials).flatten()
    I_p = t / (p_model * (1 - p_model))

    # Avoiding creating the large diagonal matrix and storing in memory
    I_w = jnp.dot((jac_full.T * I_p), jac_full) / len(p_model)
    
    # Avoiding multiplying the matrices out and calculating the trace explicitly
    return jnp.sum(jnp.multiply(jac_full.T, jnp.dot(jnp.linalg.inv(I_w), jac_full.T)))

However, when I try to call this function directly or during a jax.grad() call, I get an error message that bfloat16 is unsupported:

KeyError: ('inv', { lambda ; a:bf16[1626,1626]. let
    b:i32[1626,1626] = iota[dimension=0 dtype=int32 shape=(1626, 1626)] 
    c:i32[1626,1626] = add b 0
    d:i32[1626,1626] = iota[dimension=1 dtype=int32 shape=(1626, 1626)] 
    e:bool[1626,1626] = eq c d
    f:bf16[1626,1626] = convert_element_type[new_dtype=bfloat16 weak_type=False] e
    g:bf16[1626,1626] = xla_call[
      call_jaxpr={ lambda ; h:bf16[1626,1626] i:bf16[1626,1626]. let
          j:bf16[1626,1626] = stop_gradient h
          k:bf16[1626,1626] _:i32[1626] l:i32[1626] = lu j
          m:bf16[1626,1626] = transpose[permutation=(1, 0)] i
          n:bf16[1626,1626] = custom_linear_solve[
            const_lengths=_LinearSolveTuple(matvec=1, vecmat=1, solve=2, transpose_solve=2)
            jaxprs=_LinearSolveTuple(matvec={ lambda ; a:bf16[1626,1626] b:bf16[1626,1626]. let
    c:bf16[1626,1626] = dot_general[
      dimension_numbers=(((1,), (1,)), ((), ()))
      precision=(<Precision.HIGHEST: 2>, <Precision.HIGHEST: 2>)
      preferred_element_type=None
...
    160       np.dtype(dtype), batch, m, n)
    161   workspace = ir.RankedTensorType.get([lwork], a_type.element_type)
    162   kernel = f"{platform}solver_getrf"

ValueError: Unsupported dtype dtype(bfloat16)

I am not sure if this is a limitation of my hardware (not compatible with bfloat16?) or if the values in the array are outside of the range that bfloat16 can represent. Would appreciate any advice here!

[jakevdp]

The limitation is that the matrix inverse is not implemented in bfloat16:

import jax.numpy as jnp
x = jnp.ones((3, 3), dtype='bfloat16')
jnp.linalg.inv(x)
# NotImplementedError: Unsupported dtype bfloat16

One way you could work around this by casting to float32:

jnp.linalg.inv(x.astype('float32')).astype('bfloat16')

One aside: in general, you should avoid computing an explicit floating-point inverse, and use a linear solve instead.
So instead of jnp.dot(jnp.linalg.inv(I_w), jac_full.T), you should write:

jnp.linalg.solve(I_w, jac_full.T)

Here solve has the same issue as inv; if the arguments are bfloat16 you'll have to cast to float32.

[pvasired]

Thanks, this answers my question. Is implementing matrix inversion in bfloat16 a major challenge?

Appreciate the advice about using jnp.linalg.solve as well -- what's the reasoning behind avoiding explicit inverse calculations? I'm actually finding that changing to jnp.linalg.solve makes my code reasonably slower. For some context, I'm taking gradients of the above function and optimizing using optax

[jakevdp]

what's the reasoning behind avoiding explicit inverse calculations?

An inverse followed by a dot product accumulates more floating point roundoff error than a single numerical solve, particularly if the matrix is ill-contitioned.

Is implementing matrix inversion in bfloat16 a major challenge?

16-bit types generally aren't supported by lapack and blas libraries. For example, you'll find that numpy does not support float16 inverse either:

import numpy as np
np.linalg.inv(np.random.rand(3, 3).astype('float16'))
# TypeError: array type float16 is unsupported in linalg

[pvasired]

Thanks again!