keras/src/layers/core/einsum_dense.py/EinsumDense/grad_fn
class EinsumDense:    def _int8_call(self, inputs):
        @ops.custom_gradient
        def einsum_with_inputs_gradient(inputs, kernel, kernel_scale):
            def grad_fn(*args, upstream=None):
                if upstream is None:
                    (upstream,) = args
                # De-scale kernel
                _kernel_scale = kernel_scale  # Overcome UnboundLocalError
                if self._kernel_squeeze_axes:
                    _kernel_scale = ops.expand_dims(
                        _kernel_scale, axis=self._kernel_squeeze_axes
                    )
                if self._kernel_expand_axes:
                    _kernel_scale = ops.squeeze(
                        _kernel_scale, axis=self._kernel_expand_axes
                    )
                _kernel_scale = ops.transpose(
                    _kernel_scale, self._kernel_reverse_transpose_axes
                )
                float_kernel = ops.divide(
                    ops.cast(kernel, dtype=self.compute_dtype),
                    _kernel_scale,
                )
                # From https://stackoverflow.com/a/47609896
                inputs_grad = ops.einsum(
                    self._custom_gradient_equation, upstream, float_kernel
                )
                return (inputs_grad, None, None)

            inputs, inputs_scale = self.inputs_quantizer(inputs)
            x = ops.einsum(self.equation, inputs, kernel)
            # Deal with `inputs_scale`
            inputs_scale = ops.transpose(
                inputs_scale, self._input_transpose_axes
            )
            if self._input_expand_axes:
                inputs_scale = ops.expand_dims(
                    inputs_scale, axis=self._input_expand_axes
                )
            if self._input_squeeze_axes:
                inputs_scale = ops.squeeze(
                    inputs_scale, axis=self._input_squeeze_axes
                )
            # De-scale outputs
            x = ops.cast(x, self.compute_dtype)
            x = ops.divide(x, ops.multiply(inputs_scale, kernel_scale))
            return x, grad_fn

        x = einsum_with_inputs_gradient(
            inputs,
            ops.convert_to_tensor(self._kernel),
            ops.convert_to_tensor(self.kernel_scale),
        )
        if self.lora_enabled:
            lora_x = ops.einsum(self.equation, inputs, self.lora_kernel_a)
            lora_x = ops.matmul(lora_x, self.lora_kernel_b)
            x = ops.add(x, lora_x)
        if self.bias is not None:
            x += self.bias
        if self.activation is not None:
            x = self.activation(x)
        return x