def backward(self, inputs, output, grad, weights, state, new_state, rng=None,
               **kwargs):
    """Custom backward pass, for efficiency (see forward_and_or_backward)."""
    assert not self._use_reference_code
    del output, state, kwargs
    _, _, inputs_grad, weights_grad = self.forward_and_or_backward(
        inputs, weights, new_state, rng, output_grad=grad,
        compute_output=False, update_state=False)
    return inputs_grad, weights_grad
Custom backward pass, for efficiency (see forward_and_or_backward).
trax/layers/research/efficient_attention.py