using stax.Cos(a=1.0, b=1.0, c=0.0) to get kernel from conv layer gives error about neural-tangents HOT 2 CLOSED

Do you have a code snippet? Usually this error arises if you stack two nonlinearities one after another (e.g. Cos after ReLU), in which case the exact analytic infinite width limit is not known; but when inputs to nonlinearities are iid Gaussians (i.e. after Conv, Dense, ConvLocal, ConvTranspose etc layers with iid random weights), then the limiting covariance matrix of nonlinearity outputs can be computed.

from neural-tangents.

If there's no flatten layer, the covariance of activations, which have shapes (n1, h, w, c) and (n2, h, w, c), will tend to a dense 6D matrix of shape (n1, n2, h, h, w, w), as c is taken to infinity. This is why the covariance is 6D. When a flattened layer is used, the network output covariance has shape (n1, n2) (2D), and to compute it, only the diagonal entries of intermediary layer covariances are needed, hence it only computes a 4D (n1, n2, h, w) matrix in intermediary layers (and not full 6D covariances).

how i can reduce it to 3 or 4 dimensions

You can extract the diagonal(s) form the 6D matrix, or have your network terminate with a flattening layer. It depends on why exactly you want the shapes to be 3 or 4D.

from neural-tangents.