flax
7ee6ba63 - Add a get_inputs_kernel_bias method to the DenseGeneral layer.

Commit
2 years ago
Add a get_inputs_kernel_bias method to the DenseGeneral layer. This method returns the inputs, kernel, bias, dot_dimension_nums, and out_shape_sequence_int. This is useful for users who want to use the DenseGeneral layer with custom kernel containing the dot_general step. PiperOrigin-RevId: 629615743
Author
Committer
Parents
Loading