def backward(self, inputs, output, grad, weights, state, new_state, rng=None,
               **kwargs):
    """Custom backward pass, for efficiency (see forward_and_or_backward)."""
    del output
    del state
    del kwargs
    unused_output, unused_new_state, inputs_grad = self.forward_and_or_backward(
        inputs,
        new_state,
        rng,
        output_grad=grad,
        compute_output=False,
        update_state=False)

    weights_grad = fastmath.nested_map(np.zeros_like, weights)
    return inputs_grad, weights_grad
Custom backward pass, for efficiency (see forward_and_or_backward).
trax/layers/research/efficient_attention.py