neural_tangents.stax.FanInProd
- neural_tangents.stax.FanInProd()[source]
Layer construction function for a fan-in product layer.
This layer takes a number of inputs (e.g. produced by FanOut) and elementwise-multiplies the inputs to produce a single output.
- Return type
Tuple[InitFn,ApplyFn,LayerKernelFn,MaskFn]- Returns
(init_fn, apply_fn, kernel_fn).