Skip to content

Infer wp.array DType from argument in constructor #1539

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
thomasbbrunner wants to merge 1 commit into
base: main
Choose a base branch
from
+21 −14
Draft

Infer wp.array DType from argument in constructor #1539

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions warp/__init__.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1724,13 +1724,13 @@ class transformd:
@over
def zeros(
shape: int | tuple[int, ...] | list[int] | None = None,
dtype: type = float,
dtype: type[DType] | None = None,
device: DeviceLike = None,
requires_grad: _builtins.bool = False,
pinned: _builtins.bool = False,
retain_grad: _builtins.bool = False,
**kwargs,
) -> array:
) -> array[DType]:
"""Return a zero-initialized array."""
Comment on lines 1725 to 1734

@coderabbitai coderabbitai Bot Jun 8, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
# Verify whether DType is inferred for omitted dtype vs explicit dtype from the current stub.
# Expected:
# - reveal_type(wp.zeros(10)) should be concrete default dtype array (not Any)
# - reveal_type(wp.zeros(10, dtype=wp.float32)) should be array[float32]
cat > /tmp/verify_zeros_typing.py <<'PY'
import warp as wp

a = wp.zeros(10)
b = wp.zeros(10, dtype=wp.float32)

reveal_type(a)
reveal_type(b)
PY

# Run with your project type checker (mypy/pyright) in CI/local env.
# Example (mypy):
# mypy /tmp/verify_zeros_typing.py

Repository: NVIDIA/warp

Length of output: 37

🏁 Script executed:

#!/bin/bash
set -euo pipefail

# 1) Inspect the current zeros stubs
echo "== zeros definitions =="
rg -n "def zeros\\(" warp/__init__.pyi || true

echo
echo "== zeros snippet (around line ~1725) =="
sed -n '1700,1760p' warp/__init__.pyi

echo
echo "== dtype aliases referenced by zeros snippet =="
rg -n "class DType|DType\\b|float32\\b" warp/__init__.pyi | head -n 80

# 2) Create a reveal_type snippet
cat > /tmp/verify_zeros_typing.py <<'PY'
import warp as wp

a = wp.zeros(10)
b = wp.zeros(10, dtype=wp.float32)

reveal_type(a)
reveal_type(b)
PY

# 3) Check if type checkers exist and run them
echo
echo "== type checker availability =="
command -v mypy || true
command -v pyright || true

echo
echo "== running pyright (if available) =="
if command -v pyright >/dev/null 2>&1; then
  # Ensure pyright can see the local stub package
  pyright /tmp/verify_zeros_typing.py --pythonpath . || true
else
  echo "pyright not available"
fi

echo
echo "== running mypy (if available) =="
if command -v mypy >/dev/null 2>&1 || python -m mypy --version >/dev/null 2>&1; then
  # Make mypy find the local warp package directory (with __init__.pyi)
  MYPYPATH="$PWD" python -m mypy --hide-error-context --show-traceback /tmp/verify_zeros_typing.py || true
else
  echo "mypy not available"
fi

Repository: NVIDIA/warp

Length of output: 6199

🏁 Script executed:

#!/bin/bash
set -euo pipefail

echo "== locate zeros implementation (python) =="
rg -n "def zeros\\(" -S warp || true

echo
echo "== locate zeros callsites / dtype default =="
# Search for "zeros(" and "dtype" near it
rg -n "zeros\\(" warp | head -n 50 || true

echo
echo "== search for the core zeros implementation in _src/ =="
rg -n "zeros\\b" warp/_src | head -n 50 || true

echo
echo "== try to open likely implementation files =="
# List candidate files containing "zeros" to inspect around dtype defaults
fd -t f "*/**/*.py" warp | head -n 200 >/tmp/pyfiles.txt || true
python3 - <<'PY'
import subprocess, re, os, pathlib, sys
root="warp"
# Use ripgrep to find files with 'zeros' keyword
cmd=["rg","-l","\\bzeros\\b",root]
out=subprocess.check_output(cmd, text=True).splitlines()
for p in out[:25]:
    print(p)
PY

echo
echo "== show first matches around zeros in candidate files =="
# Show contexts for zeros definitions if any
for f in $(rg -l "def zeros\\(" warp || true); do
  echo "--- $f ---"
  rg -n "def zeros\\(" "$f" -n -C 3 || true
done

echo
echo "== show contexts for 'dtype' default in zeros implementations =="
# Broad heuristic: within files mentioning zeros, show nearby 'dtype' and 'float' references
for f in $(rg -l "\\bzeros\\b" warp || true); do
  if rg -q "zeros\\b" "$f"; then
    # Only display if file is reasonably small for context scanning
    lines=$(wc -l < "$f")
    if [ "$lines" -le 600 ]; then
      echo "--- $f ---"
      rg -n "zeros\\b|dtype\\b|float\\b" "$f" -C 2 || true
    fi
  fi
done

Repository: NVIDIA/warp

Length of output: 50675

🏁 Script executed:

#!/bin/bash
set -euo pipefail

# show zeros implementation signature + dtype default handling
sed -n '7635,7705p' warp/_src/context.py

echo "---- dtype handling (search within function) ----"
python3 - <<'PY'
import re, pathlib
p = pathlib.Path("warp/_src/context.py")
text = p.read_text()
# naive slice around function def
start = text.find("def zeros(")
print("function start idx", start)
sub = text[start:start+6000]
# print lines containing "dtype" and likely defaults
for i,line in enumerate(sub.splitlines(),1):
    if "dtype" in line or "float" in line or "DType" in line or "float32" in line:
        if any(k in line for k in ["dtype", "float32", "float", "default", "None", "type"]):
            print(f"{i:04d}: {line}")
PY

echo "---- exact dtype default line(s) ----"
rg -n "def zeros\\(|dtype\\s*=|if\\s+dtype\\s+is\\s+None|dtype\\s+is\\s+None|float32|float\\b" -n warp/_src/context.py | head -n 60

Repository: NVIDIA/warp

Length of output: 7298

🏁 Script executed:

#!/bin/bash
set -euo pipefail

echo "== header: find float/float32 dtype aliases =="
sed -n '1,120p' warp/__init__.pyi | sed -n '1,120p'

echo
echo "== locate DType definition area =="
sed -n '230,320p' warp/__init__.pyi

echo
echo "== locate array generic + dtype typing =="
rg -n "class Array\\b|class FabricArray\\b|DType\\b\\s*=\\s*TypeVar" warp/__init__.pyi | head -n 50

echo
echo "== show around warp.float / float32 declarations =="
rg -n "(^|\\b)(float32|float\\b|int32|int\\b)\\s*=" warp/__init__.pyi | head -n 80

echo
echo "== show type alias for float if present =="
# Try to find "float =" or "float: "
rg -n "^[[:space:]]*float[[:space:]]*=" warp/__init__.pyi || true
rg -n "^[[:space:]]*float\\b.*:" warp/__init__.pyi || true

echo
echo "== show Array[...] return typing helper =="
rg -n "def zeros\\(" warp/__init__.pyi | head -n 20

Repository: NVIDIA/warp

Length of output: 10632

🏁 Script executed:

#!/bin/bash
set -euo pipefail

echo "== zeros related empty stub =="
rg -n "def empty\\(" warp/__init__.pyi | head -n 20

# show area around empty signature
line=$(rg -n "def empty\\(" warp/__init__.pyi | head -n 1 | cut -d: -f1)
start=$((line-10))
end=$((line+60))
sed -n "${start},${end}p" warp/__init__.pyi

echo
echo "== zeros_like/ones default dtype typing (if present) =="
rg -n "def (ones|full|empty_like|zeros_like|ones_like|full_like)\\(" warp/__init__.pyi | head -n 50

Repository: NVIDIA/warp

Length of output: 86

🏁 Script executed:

#!/bin/bash
set -euo pipefail

echo "== stub block: warp/__init__.pyi zeros overloads =="
sed -n '1715,1760p' warp/__init__.pyi

echo
echo "== runtime: warp/_src/context.py empty() dtype mapping for python float =="
# locate empty() definition
rg -n "def empty\\(" warp/_src/context.py | head -n 20
line_empty=$(rg -n "def empty\\(" warp/_src/context.py | head -n 1 | cut -d: -f1)
start=$((line_empty-5))
end=$((line_empty+120))
sed -n "${start},${end}p" warp/_src/context.py

echo
echo "== runtime: dtype=float mapping if dtype is None => float (zeros) => empty() =="
rg -n "if dtype is None:|dtype = float|if dtype is float|dtype == float" warp/_src/context.py | head -n 80

echo
echo "== locate typing Scalar definition used by __init__.pyi =="
# Scalar is imported from warp._src.types; search for its definition in that file
rg -n "^(Scalar\\s*=|class Scalar\\b|from typing import .*Scalar)" warp/_src/types.py warp/_src/types.pyi 2>/dev/null || true
rg -n "\\bScalar\\b\\s*=" warp/_src/types.py warp/_src/types.pyi 2>/dev/null || true

echo
echo "== show relevant snippet around Scalar (small window) =="
# If a Scalar definition exists in types.py, print around the first match
scalar_line=$(rg -n "Scalar\\b\\s*=" warp/_src/types.py 2>/dev/null | head -n 1 | cut -d: -f1 || true)
if [ -n "${scalar_line:-}" ]; then
  start=$((scalar_line-20))
  end=$((scalar_line+60))
  sed -n "${start},${end}p" warp/_src/types.py
else
  echo "No direct 'Scalar=' found in warp/_src/types.py; searching for 'Scalar' typing constructs..."
  rg -n "Scalar\\b" warp/_src/types.py | head -n 40
fi

Repository: NVIDIA/warp

Length of output: 8692

🏁 Script executed:

#!/bin/bash
set -euo pipefail

echo "== locate type_to_warp mapping =="
rg -n "def type_to_warp\\b|type_to_warp\\b" warp/_src/types.py

echo
echo "== show mapping logic around type_to_warp =="
line=$(rg -n "def type_to_warp\\b" warp/_src/types.py | head -n 1 | cut -d: -f1)
start=$((line-30))
end=$((line+120))
sed -n "${start},${end}p" warp/_src/types.py

echo
echo "== look for python float -> float32 mapping explicitly =="
rg -n "float32\\b.*(float\\b|python float)|\\bfloat\\b.*float32\\b" warp/_src/types.py | head -n 50

echo
echo "== show native_scalar_types and any mapping usage =="
rg -n "native_scalar_types|float32\\b|python\\b.*float" warp/_src/types.py | head -n 80

Repository: NVIDIA/warp

Length of output: 11202

Fix zeros() stub to bind return dtype when dtype is omitted

warp/_src/context.py::zeros() defaults dtype to Python float when omitted, and warp/_src/types.py::type_to_warp() maps that to warp.float32; however the current stub keeps dtype as type[DType] | None = None, so wp.zeros(10) can’t bind DType and the return often falls back to an imprecise type (e.g., array[Any]/unknown). Split overloads so omitted dtype returns array[float32] and explicit dtype binds DType.

Suggested stub shape 
 `@over`
 def zeros(
     shape: int | tuple[int, ...] | list[int] | None = None,
-    dtype: type[DType] | None = None,
+    dtype: None = None,
     device: DeviceLike = None,
     requires_grad: _builtins.bool = False,
     pinned: _builtins.bool = False,
     retain_grad: _builtins.bool = False,
     **kwargs,
-) -> array[DType]:
+) -> array[float32]: ...
+
+@over
+def zeros(
+    shape: int | tuple[int, ...] | list[int] | None = None,
+    dtype: type[DType] = ...,
+    device: DeviceLike = None,
+    requires_grad: _builtins.bool = False,
+    pinned: _builtins.bool = False,
+    retain_grad: _builtins.bool = False,
+    **kwargs,
+) -> array[DType]:
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
def zeros(
shape: int | tuple[int, ...] | list[int] | None = None,
dtype: type = float,
dtype: type[DType] | None = None,
device: DeviceLike = None,
requires_grad: _builtins.bool = False,
pinned: _builtins.bool = False,
retain_grad: _builtins.bool = False,
**kwargs,
) -> array:
) -> array[DType]:
"""Return a zero-initialized array."""
`@over`
def zeros(
shape: int | tuple[int, ...] | list[int] | None = None,
dtype: None = None,
device: DeviceLike = None,
requires_grad: _builtins.bool = False,
pinned: _builtins.bool = False,
retain_grad: _builtins.bool = False,
**kwargs,
) -> array[float32]: ...
`@over`
def zeros(
shape: int | tuple[int, ...] | list[int] | None = None,
dtype: type[DType] = ...,
device: DeviceLike = None,
requires_grad: _builtins.bool = False,
pinned: _builtins.bool = False,
retain_grad: _builtins.bool = False,
**kwargs,
) -> array[DType]:
"""Return a zero-initialized array."""
🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@warp/__init__.pyi` around lines 1725 - 1734, The zeros() stub must be split
into two overloads so the return dtype binds when dtype is provided and defaults
to warp.float32 when omitted: add one overload for zeros(..., dtype: None =
None) -> array[float32] (or the module's float32 DType symbol) and a generic
overload zeros(..., dtype: type[DType]) -> array[DType]; keep other params
(shape, device, requires_grad, pinned, retain_grad, **kwargs) identical and
ensure the TypeVar DType and array type are the same symbols used elsewhere in
the stub to allow proper type inference for zeros(... ) and zeros(...,
dtype=...).

...

Expand Down
29 changes: 18 additions & 11 deletions warp/_src/context.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -66,7 +66,7 @@ def ParamSpec(name):
from warp._src.codegen import WarpCodegenTypeError, synchronized
from warp._src.logger import LOG_DEBUG, LOG_WARNING, get_logger, log_debug, log_error, log_info, log_warning
from warp._src.texture import Texture1D, Texture2D, Texture3D, texture1d_t, texture2d_t, texture3d_t
from warp._src.types import LAUNCH_MAX_DIMS, Array, LaunchBounds, launch_bounds_t, type_repr
from warp._src.types import LAUNCH_MAX_DIMS, Array, DType, LaunchBounds, launch_bounds_t, type_repr

_wp_module_name_ = "warp.context"

Expand Down Expand Up @@ -7647,13 +7647,13 @@ def unmap(self):

def zeros(
shape: int | tuple[int, ...] | list[int] | None = None,
dtype: type = float,
dtype: type[DType] | None = None,

@thomasbbrunner thomasbbrunner Jun 8, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The default was changed to None.

This is to avoid a type error, because the concrete default float would pin the free DType TypeVar.

There is precedent to this (wp.full already uses this pattern).

This changes the behavior of the code. wp.zeros(10, dtype=None) is now a legal "use the default". Before it raised a TypeError in type_size_in_bytes.

@christophercrouzet christophercrouzet Jun 9, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's better to have float set as the default parameter value since the expectations are clear from the signature but, other than that, I don't think it's too much of an issue since it's a common Python idiom to consider None as a way to express “no value has been passed” and define the default behaviour inside the function's body.

device: DeviceLike = None,
requires_grad: bool = False,
pinned: bool = False,
retain_grad: bool = False,
**kwargs,
) -> warp.array:
) -> warp.array[DType]:
Comment on lines 7648 to +7656

@greptile-apps greptile-apps Bot Jun 8, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Default dtype not statically reflected when omitted

Changing the parameter default from dtype: type = float to dtype: type[DType] | None = None means type checkers bind DType only when a concrete dtype is passed. When the user calls wp.zeros(10) (omitting dtype), DType remains unresolved and the return type degrades to array[Any] — the same situation mypy strict mode complained about before. The intended default of float / wp.float32 is only applied at runtime via the if dtype is None guard, invisible to the type checker.

Adding @typing.overload variants (one for dtype=None returning array[float32] and one for dtype: type[DType] returning array[DType]) would give fully typed results in both cases. This pattern is also needed for ones and empty, which receive the same treatment here.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

"""Return a zero-initialized array.

Args:
Expand All @@ -7667,6 +7667,8 @@ def zeros(
Returns:
A warp.array object representing the allocation
"""
if dtype is None:
dtype = float

arr = empty(
shape=shape,
Expand Down Expand Up @@ -7712,13 +7714,13 @@ def zeros_like(

def ones(
shape: int | tuple[int, ...] | list[int] | None = None,
dtype: type = float,
dtype: type[DType] | None = None,
device: DeviceLike = None,
requires_grad: bool = False,
pinned: bool = False,
retain_grad: bool = False,
**kwargs,
) -> warp.array:
) -> warp.array[DType]:
"""Return a one-initialized array.

Args:
Expand All @@ -7732,6 +7734,8 @@ def ones(
Returns:
A warp.array object representing the allocation
"""
if dtype is None:
dtype = float

return full(
shape=shape,
Expand Down Expand Up @@ -7771,13 +7775,13 @@ def ones_like(
def full(
shape: int | tuple[int, ...] | list[int] | None = None,
value: Any = 0,
dtype: type | None = None,
dtype: type[DType] | None = None,
device: DeviceLike = None,
requires_grad: bool = False,
pinned: bool = False,
retain_grad: bool = False,
**kwargs,
) -> warp.array:
) -> warp.array[DType]:
"""Return an array with all elements initialized to the given value.

Args:
Expand Down Expand Up @@ -7902,13 +7906,13 @@ def clone(

def empty(
shape: int | tuple[int, ...] | list[int] | None = None,
dtype=float,
dtype: type[DType] | None = None,
device: DeviceLike = None,
requires_grad: bool = False,
pinned: bool = False,
retain_grad: bool = False,
**kwargs,
) -> warp.array:
) -> warp.array[DType]:
"""Return an uninitialized array.

Args:
Expand All @@ -7923,6 +7927,9 @@ def empty(
A warp.array object representing the allocation
"""

if dtype is None:
dtype = float

# backwards compatibility for case where users called wp.empty(n=length, ...)
if "n" in kwargs:
shape = (kwargs["n"],)
Expand Down Expand Up @@ -7998,12 +8005,12 @@ def empty_like(

def from_numpy(
arr: np.ndarray,
dtype: type | None = None,
dtype: type[DType] | None = None,
shape: Sequence[int] | None = None,
device: DeviceLike | None = None,
requires_grad: bool = False,
retain_grad: bool = False,
) -> warp.array:
) -> warp.array[DType]:
Comment on lines 8006 to +8013

@greptile-apps greptile-apps Bot Jun 8, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 from_numpy with omitted dtype still returns array[Any]

When dtype is not supplied, the runtime successfully infers it from the numpy array's .dtype (resolving to e.g. float32, vec3, etc.), but the type checker sees an unbound DType and resolves the return as array[Any]. There is no static way to bridge the numpy dtype → Warp dtype mapping without explicit @overloads or a different approach — so this is a known limitation — but it's worth noting that the annotation currently overpromises for the no-dtype path: callers relying on inferred types from from_numpy(arr) won't benefit from this PR.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

"""Return a Warp array created from a NumPy array.

Args:
Expand Down
2 changes: 1 addition & 1 deletion warp/_src/types.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3072,7 +3072,7 @@ def __new__(cls, *args, **kwargs):
def __init__(
self,
data: list | tuple | npt.NDArray | None = None,
dtype: Any = Any,
dtype: type[DType] = Any,
shape: int | tuple[int, ...] | list[int] | None = None,
strides: tuple[int, ...] | None = None,
ptr: int | None = None,
Expand Down