MIL pow op for 2.0 to 0.0 is 0.0

[kentslaney]

🐞Describing the bug

mb.pow(x=2.0, y=0.0) != 1.0, a discrepancy that isn't mentioned in the API reference.

To Reproduce

import numpy as np
import coremltools as ct
from coremltools.converters.mil.mil import Builder as mb
from coremltools.converters.mil.mil import types

def mb_pow_issue():
    @mb.program(input_specs=[mb.TensorSpec(shape=(1,), dtype=types.fp32)])
    def prog(e_raw):
        # Problematic approach: mb.pow(2, e_raw)
        # We cast 2 to the same type as e_raw (float32)
        two = mb.cast(x=2.0, dtype="fp32")
        pow_2_e_raw_unstable = mb.pow(x=two, y=e_raw, name="unstable_output")
        
        # Workaround: exp(e_raw * ln(2))
        ln_2 = np.log(2.0).astype(np.float32)
        pow_2_e_raw_stable = mb.exp(x=mb.mul(x=e_raw, y=ln_2), name="stable_output")
        
        return pow_2_e_raw_unstable, pow_2_e_raw_stable

    # Compile the model
    model = ct.convert(prog, source='milinternal', convert_to="mlprogram", compute_precision=ct.precision.FLOAT32)

    # Test with 0
    e_raw_val = np.array([0.0], dtype=np.float32)
    prediction = model.predict({"e_raw": e_raw_val})
    
    print(f"Input e_raw: {e_raw_val}")
    print(f"Unstable output (mb.pow): {prediction['unstable_output']}")
    print(f"Stable output (exp trick): {prediction['stable_output']}")

if __name__ == "__main__":
    mb_pow_issue()

which prints

Input e_raw: [0.]
Unstable output (mb.pow): [0.]
Stable output (exp trick): [1.]

System environment (please complete the following information):

• coremltools version: 9.0
• OS (e.g. MacOS version or Linux type): macOS 26.1

[TobyRoseman]

Something strange is happening here .

Using

pow_2_e_raw_unstable = mb.pow(x=two, y=0.0, name="unstable_output")

Gives the correct results.

Changing the return statement to:

return pow_2_e_raw_unstable, pow_2_e_raw_stable, e_raw

Also gives the correct result, and e_raw is 0..

It doesn't seem related to the compute precision or any optimization passes. With the original prog, the following conversion still has the issue:

model = ct.convert(prog, source='milinternal', convert_to="mlprogram",  pass_pipeline=ct.PassPipeline.EMPTY)

Even stranger, the following code works:

import numpy as np
import coremltools as ct
from coremltools.converters.mil.mil import Builder as mb
from coremltools.converters.mil.mil import types

@mb.program(input_specs=[mb.TensorSpec(shape=(1,), dtype=types.fp32)])
def prog(e_raw):
  two = mb.cast(x=2.0, dtype="fp32")
  pow_2_e_raw_unstable = mb.pow(x=two, y=e_raw, name="unstable_output")
        
  return pow_2_e_raw_unstable

model = ct.convert(prog, source='milinternal', convert_to="mlprogram", compute_precision=ct.precision.FLOAT32)
e_raw_val = np.array([0.0], dtype=np.float32)
print(model.predict({"e_raw": e_raw_val}))

It prints: {'unstable_output': array([1.], dtype=float32)}

[kasper0406]

I also had a quick look into this, without making any progress. Just noting my findings, which aligns with what @TobyRoseman describes.

The MIL program generated for the cases of not returning e_raw (not working), and returning e_raw (working) is the following:

# Not returning e_raw
main[CoreML3](%e_raw: (1,fp32)(Tensor)) {
  block0() {
    %cast_0: (fp32)*(Scalar) = cast(x=2.0, dtype="fp32", name="cast_0")
    %unstable_output: (1,fp32)(Tensor) = pow(x=%cast_0, y=%e_raw, name="unstable_output")
    %mul_0: (1,fp32)(Tensor) = mul(x=%e_raw, y=0.6931471824645996, name="mul_0")
    %stable_output: (1,fp32)(Tensor) = exp(x=%mul_0, name="stable_output")
  } -> (%unstable_output, %stable_output)
}

# Returning e_raw
main[CoreML3](%e_raw: (1,fp32)(Tensor)) {
  block0() {
    %cast_0: (fp32)*(Scalar) = cast(x=2.0, dtype="fp32", name="cast_0")
    %unstable_output: (1,fp32)(Tensor) = pow(x=%cast_0, y=%e_raw, name="unstable_output")
    %mul_0: (1,fp32)(Tensor) = mul(x=%e_raw, y=0.6931471824645996, name="mul_0")
    %stable_output: (1,fp32)(Tensor) = exp(x=%mul_0, name="stable_output")
  } -> (%unstable_output, %stable_output, %e_raw)
}

So at least what I can tell from the MIL side, there is nothing wrong with the programs.

I tried saving the not working (i.e. the one not returning e_raw) model using model.save and calling it from Swift, to rule out potential issues in the coremltools model-calling logic:

import CoreML

let config = MLModelConfiguration()
config.computeUnits = .cpuOnly
let model = try test_pow(configuration: config)

let e_raw_val = try MLMultiArray.init(shape: [1], dataType: .float32)
e_raw_val[0] = 0.0

let predictions = try model.prediction(e_raw: e_raw_val)

print("Stable output \(predictions.stable_output)")
print("Unstable output \(predictions.unstable_output)")

From Swift it is still producing unexpected output:

Stable output Float32 1 vector
[1]
Unstable output Float32 1 vector
[0]

[kentslaney]

Running _get_mil_internal on the converted model gives

main[CoreML3](%e_raw: (1,fp32)(Tensor)) {
  block0() {
    %unstable_output: (1,fp32)(Tensor) = pow(x=2.0, y=%e_raw, name="unstable_output")
    %mul_0: (1,fp32)(Tensor) = mul(x=%e_raw, y=0.6931471824645996, name="mul_0")
    %stable_output: (1,fp32)(Tensor) = exp(x=%mul_0, name="stable_output")
  } -> (%unstable_output, %stable_output)
}

[TobyRoseman]

Just printing the MIL program is probably not telling the whole story. IIRC - certain ops, such as constant ops, are not shown when the MIL program is printed.