def size(x):
x = x.numpy() if isinstance(x, tf.Tensor) else x
if isinstance(
x, (int, float, onp.int32, onp.int64, onp.float32, onp.float64)):
return 1
elif isinstance(x, (np.ndarray, onp.ndarray)):
return np.prod(x.shape)
else:
raise TypeError(
"The inputs must be one of types {int, float, numpy array"
", TensorFlow Tensor, TensorFlow ndarray} object.")
def reduce_window(inputs, init_value, reducer, window_dimensions, strides,
padding, base_dilation=None, window_dilation=None):
# Add an extra "batch" dimension and an extra "channel" dimension to pass the
# TensorFlow pool dimensionality checker.
inputs = np.expand_dims(inputs, axis=(0, inputs.ndim))
if reducer not in [np.max, np.add]:
raise TypeError('Only max pooling and average/sum pooling are supported.')
# Note that there is no need to send in the parameter data format since the
# input is already of default data format - 'N...C'. The adjustments of the
# input shape is already finished in apply_fun of Pooling in stax.
pooling = 'AVG' if reducer == np.add else 'MAX'
output = np.asarray(nn.pool(inputs, window_dimensions, pooling, strides, padding))
return np.squeeze(output, axis=(0, output.ndim - 1)) * np.prod(window_dimensions)
def run_test(*args):
num_samples = 1000
tol = 0.1 # High tolerance to keep the # of samples low else the test
# takes a long time to run.
np_random.seed(10)
outputs = [np_random.randn(*args) for _ in range(num_samples)]
# Test output shape.
for output in outputs:
self.assertEqual(output.shape, tuple(args))
default_dtype = (np.float64 if np_dtypes.is_allow_float64()
else np.float32)
self.assertEqual(output.dtype.as_numpy_dtype, default_dtype)
if np.prod(args): # Don't bother with empty arrays.
outputs = [output.tolist() for output in outputs]
# Test that the properties of normal distribution are satisfied.
mean = np.mean(outputs, axis=0)
stddev = np.std(outputs, axis=0)
self.assertAllClose(mean, np.zeros(args), atol=tol)
self.assertAllClose(stddev, np.ones(args), atol=tol)
# Test that outputs are different with different seeds.
np_random.seed(20)
diff_seed_outputs = [
np_random.randn(*args).tolist() for _ in range(num_samples)
]
self.assertNotAllClose(outputs, diff_seed_outputs)
# Test that outputs are the same with the same seed.
np_random.seed(10)
same_seed_outputs = [
np_random.randn(*args).tolist() for _ in range(num_samples)
]
self.assertAllClose(outputs, same_seed_outputs)