mirror of
https://github.com/saymrwulf/pytorch.git
synced 2026-05-15 21:00:47 +00:00
9945fd7253
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/49980 From ``` ./python/libcst/libcst codemod remove_unused_imports.RemoveUnusedImportsWithGlean --no-format caffe2/ ``` Test Plan: Standard sandcastle tests Reviewed By: xush6528 Differential Revision: D25727359 fbshipit-source-id: c4f60005b10546423dc093d31d46deb418352286
53 lines
1.7 KiB
Python
53 lines
1.7 KiB
Python
|
|
|
|
|
|
|
|
|
|
from caffe2.python import core
|
|
|
|
import caffe2.python.hypothesis_test_util as hu
|
|
import caffe2.python.serialized_test.serialized_test_util as serial
|
|
import hypothesis.strategies as st
|
|
import itertools as it
|
|
import numpy as np
|
|
|
|
|
|
class TestMomentsOp(serial.SerializedTestCase):
|
|
def run_moments_test(self, X, axes, keepdims, gc, dc):
|
|
if axes is None:
|
|
op = core.CreateOperator(
|
|
"Moments",
|
|
["X"],
|
|
["mean", "variance"],
|
|
keepdims=keepdims,
|
|
)
|
|
else:
|
|
op = core.CreateOperator(
|
|
"Moments",
|
|
["X"],
|
|
["mean", "variance"],
|
|
axes=axes,
|
|
keepdims=keepdims,
|
|
)
|
|
|
|
def ref(X):
|
|
mean = np.mean(X, axis=None if axes is None else tuple(
|
|
axes), keepdims=keepdims)
|
|
variance = np.var(X, axis=None if axes is None else tuple(
|
|
axes), keepdims=keepdims)
|
|
return [mean, variance]
|
|
|
|
self.assertReferenceChecks(gc, op, [X], ref)
|
|
self.assertDeviceChecks(dc, op, [X], [0, 1])
|
|
self.assertGradientChecks(gc, op, [X], 0, [0, 1])
|
|
|
|
@serial.given(X=hu.tensor(dtype=np.float32), keepdims=st.booleans(),
|
|
num_axes=st.integers(1, 4), **hu.gcs)
|
|
def test_moments(self, X, keepdims, num_axes, gc, dc):
|
|
self.run_moments_test(X, None, keepdims, gc, dc)
|
|
num_dims = len(X.shape)
|
|
if num_dims < num_axes:
|
|
self.run_moments_test(X, range(num_dims), keepdims, gc, dc)
|
|
else:
|
|
for axes in it.combinations(range(num_dims), num_axes):
|
|
self.run_moments_test(X, axes, keepdims, gc, dc)
|