pytorch/caffe2/python/operator_test/pow_op_test.py
Aapo Kyrola d8588d8007 CUDA version of elementwise power + rename to Pow + gradient
Summary: Renamed ElementwisePower to Pow for better discoverability. Added CUDA version and Gradient + tests.

Reviewed By: kennyhorror

Differential Revision: D4665550

fbshipit-source-id: dd33d8ad3917d71504e363ab397af50d38a63b1f
2017-03-07 10:20:40 -08:00

35 lines
1,017 B
Python

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals
from caffe2.python import core
from hypothesis import given
from hypothesis import strategies as st
import caffe2.python.hypothesis_test_util as hu
import unittest
class TestPowOp(hu.HypothesisTestCase):
@given(X=hu.tensor(),
exponent=st.floats(min_value=2.0, max_value=3.0),
**hu.gcs)
def test_elementwise_power(self, X, exponent, gc, dc):
def powf(X):
return (X ** exponent,)
def powf_grad(g_out, outputs, fwd_inputs):
return (exponent * (fwd_inputs[0] ** (exponent - 1)) * g_out,)
op = core.CreateOperator(
"Pow", ["X"], ["Y"], exponent=exponent)
self.assertReferenceChecks(gc, op, [X], powf,
output_to_grad="Y",
grad_reference=powf_grad),
if __name__ == "__main__":
unittest.main()