Shape of the LinearOperator, equal to [b0, ..., bn, h, w], where
b0...bn are batch dimensions h and w are the height and width of the
matrix represented by the LinearOperator.
scale_initializer
Variables are initialized with samples from
Normal(0, scale_initializer).