Commit7fcfaa7

committed

Bring back the optimization of integer.pow({2.0, 3.0})

They were accidentally removed in#26020[ghstack-poisoned]

1 parent755b7e4 commit7fcfaa7Copy full SHA for 7fcfaa7

File tree

+13

-1

lines changed

+13

-1

lines changed

Lines changed: 13 additions & 1 deletion

Original file line number	Diff line number	Diff line change
`@@ -130,7 +130,19 @@ void pow_tensor_scalar_kernel(TensorIterator& iter, Scalar exp_scalar) {`
`130`	`130`	`// tensors to float exponent e.g. tensor([4]).pow(0.5) will be tensor([2])`
`131`	`131`	`constauto exp = exp_scalar.to<double>();`
`132`	`132`	`AT_DISPATCH_INTEGRAL_TYPES(iter.dtype(),"pow", [&]() {`
`133`		`-if (exp ==0.5) {`
	`133`	`+if (exp ==2) {`
	`134`	`+cpu_kernel(iter,`
	`135`	`+ [](scalar_t base) ->scalar_t {`
	`136`	`+return base * base;`
	`137`	`+ }`
	`138`	`+ );`
	`139`	`+ }elseif (exp ==3) {`
	`140`	`+cpu_kernel(iter,`
	`141`	`+ [](scalar_t base) ->scalar_t {`
	`142`	`+return base * base * base;`
	`143`	`+ }`
	`144`	`+ );`
	`145`	`+ }elseif (exp ==0.5) {`
`134`	`146`	`cpu_kernel(iter,`
`135`	`147`	`[](scalar_t base) ->scalar_t {`
`136`	`148`	`returnstd::sqrt(static_cast<longdouble>(base));`

Comments

(0)