Browse Source

fix(mge/functional): int operations(div, exp, pow)

GitOrigin-RevId: dc434fa7ec
tags/v1.0.0-rc1
Megvii Engine Team 4 years ago
parent
commit
56cb5d6a16
2 changed files with 13 additions and 0 deletions
  1. +9
    -0
      imperative/python/megengine/core/tensor/tensor_wrapper.py
  2. +4
    -0
      imperative/python/megengine/functional/elemwise.py

+ 9
- 0
imperative/python/megengine/core/tensor/tensor_wrapper.py View File

@@ -23,6 +23,15 @@ from .tensor import Tensor

def _elwise(*args, mode):
op = builtin.Elemwise(mode=mode)
if mode in ("TRUE_DIV", "POW"):
args = tuple(
map(
lambda x: x.astype("float32")
if hasattr(x, "dtype") and x.dtype != np.float32
else x,
args,
)
)
args = utils.convert_inputs(*args)
(result,) = apply(op, *args)
return result


+ 4
- 0
imperative/python/megengine/functional/elemwise.py View File

@@ -76,6 +76,10 @@ __all__ = [

def _elwise(*args, mode):
op = builtin.Elemwise(mode=mode)
if mode in ("true_div", "exp", "pow", "log", "expm1", "log1p"):
args = tuple(
map(lambda x: x.astype("float32") if hasattr(x, "dtype") else x, args)
)
args = utils.convert_inputs(*args)
(result,) = apply(op, *args)
return result


Loading…
Cancel
Save