Initial import: grid-bot — grid trading bot for BTC-USDT on Cifra Markets

kronos-venv/lib/python3.12/site-packages/torch/distributed/_ops/__init__.py

@@ -0,0 +1 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates

kronos-venv/lib/python3.12/site-packages/torch/distributed/_ops/device_mesh.py

@@ -0,0 +1,105 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates
+
+import torch
+from torch import SymInt
+
+from ..device_mesh import DeviceMesh
+
+
+# Register custom operator
+torch.library.define(
+    "device_mesh::_runtime_compute_coordinate_on_dim",
+    "(Tensor full_mesh, int index) -> SymInt",
+    tags=torch.Tag.pt2_compliant_tag,
+)
+
+
+@torch.library.register_fake("device_mesh::_runtime_compute_coordinate_on_dim")
+def _runtime_compute_coordinate_on_dim_fake(
+    full_mesh: torch.Tensor, index: int
+) -> SymInt:
+    from torch.fx.experimental.symbolic_shapes import _constrain_range_for_size
+
+    ctx = torch._custom_op.impl.get_ctx()
+    shape_env = ctx._shape_env
+
+    # Bypass allow_dynamic_output_shape_ops check by directly creating the symint.
+    # This is intentional - the coordinate is always valid and bounded.
+    sz = shape_env.create_unbacked_symint()
+
+    # Apply size constraints - coordinate is bounded by mesh size on the given dimension.
+    # The full_mesh tensor has an extra batch dimension at the front, so the actual
+    # mesh dimensions start at index 1. mesh.size(index) = full_mesh.size(index + 1)
+    mesh_size = full_mesh.size(index + 1)
+    _constrain_range_for_size(
+        sz, min=0, max=mesh_size - 1 if isinstance(mesh_size, int) else None
+    )
+
+    try:
+        # Check if we're currently tracing in dynamo (as opposed to AOT or export).
+        in_dynamo = torch._dynamo.symbolic_convert.InstructionTranslator.current_tx()
+    except AttributeError:
+        in_dynamo = False
+
+    if in_dynamo:
+        # During dynamo tracing, distributed ops are treated as atomic - so the
+        # rank SymInt may be computed but not traced into the graph (e.g., it
+        # affects tensor values but not shapes). Mark it as ignorable here;
+        # when we decompose these ops later (after dynamo), we'll create fresh
+        # SymInts that do get traced.
+        shape_env.ignorable_fresh_unbacked_symbols.append(sz.node._expr)
+
+    return sz
+
+
+@torch.library.impl(
+    "device_mesh::_runtime_compute_coordinate_on_dim", "CompositeExplicitAutograd"
+)
+def _runtime_compute_coordinate_on_dim_impl(full_mesh: torch.Tensor, index: int) -> int:
+    rank = torch.distributed.get_rank()
+    mesh = DeviceMesh._get_mesh_tensor_from_full_mesh(full_mesh)
+    mesh_coords = DeviceMesh._compute_coordinates_from_mesh(mesh, rank)
+    if mesh_coords is None:
+        raise AssertionError
+    return mesh_coords[index]
+
+
+def _get_flattened_submesh_impl(mesh: DeviceMesh, mesh_dims: list[int]) -> DeviceMesh:
+    from torch.distributed.tensor._redistribute import (
+        _get_flattened_mesh_by_layout_impl,
+    )
+
+    result = _get_flattened_mesh_by_layout_impl(mesh, tuple(mesh_dims))
+    if result is None:
+        raise ValueError(f"No flattened mesh found for mesh_dims={mesh_dims} on {mesh}")
+    return result
+
+
+@torch.library.custom_op("device_mesh::_get_flattened_submesh", mutates_args=())
+def _get_flattened_submesh(mesh: DeviceMesh, mesh_dims: list[int]) -> DeviceMesh:
+    return _get_flattened_submesh_impl(mesh, mesh_dims)
+
+
+@_get_flattened_submesh.register_fake
+def _get_flattened_submesh_fake(mesh: DeviceMesh, mesh_dims: list[int]) -> DeviceMesh:
+    return _get_flattened_submesh_impl(mesh, mesh_dims)
+
+
+def _get_submesh_impl(mesh: DeviceMesh, mesh_dims: list[int]) -> DeviceMesh:
+    all_dim_names = mesh._mesh_dim_names
+    if all_dim_names is None:
+        raise ValueError(f"Cannot slice mesh without dim names: {mesh}")
+    dim_names = tuple(all_dim_names[i] for i in mesh_dims)
+    if len(dim_names) == 1:
+        return mesh[dim_names[0]]
+    return mesh[dim_names]
+
+
+@torch.library.custom_op("device_mesh::_get_submesh", mutates_args=())
+def _get_submesh(mesh: DeviceMesh, mesh_dims: list[int]) -> DeviceMesh:
+    return _get_submesh_impl(mesh, mesh_dims)
+
+
+@_get_submesh.register_fake
+def _get_submesh_fake(mesh: DeviceMesh, mesh_dims: list[int]) -> DeviceMesh:
+    return _get_submesh_impl(mesh, mesh_dims)