Add back compatibility with torch

bittensor/chain_data.py

@@ -15,7 +15,7 @@
 # OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
 # DEALINGS IN THE SOFTWARE.
 import bittensor
-
+import os
 import json
 from enum import Enum
 from dataclasses import dataclass, asdict
@@ -25,7 +25,7 @@
 from scalecodec.type_registry import load_type_registry_preset
 from scalecodec.utils.ss58 import ss58_encode
 
-from .utils import networking as net, U16_MAX, U16_NORMALIZED_FLOAT
+from .utils import networking as net, U16_MAX, U16_NORMALIZED_FLOAT, maybe_get_torch
 from .utils.balance import Balance
 
 custom_rpc_type_registry = {
@@ -264,15 +264,43 @@ def from_neuron_info(cls, neuron_info: dict) -> "AxonInfo":
             coldkey=neuron_info["coldkey"],
         )
 
-    def to_parameter_dict(self) -> dict[str, Union[int, str]]:
-        r"""Returns a dict of the subnet info."""
+    def _to_parameter_dict_torch(self) -> "torch.nn.ParameterDict":
+        """Returns a torch tensor of the subnet info."""
+        return maybe_get_torch().nn.ParameterDict(self.__dict__)
+
+    def _to_parameter_dict_numpy(self) -> dict[str, Union[int, str]]:
+        """Returns a dict of the subnet info."""
         return self.__dict__
 
+    def to_parameter_dict(
+        self,
+    ) -> Union[dict[str, Union[int, str]], "torch.nn.ParameterDict"]:
+        if os.environ.get("USE_TORCH"):
+            return self._to_parameter_dict_torch()
+        else:
+            return self._to_parameter_dict_numpy()
+
     @classmethod
-    def from_parameter_dict(cls, parameter_dict: dict[str, Any]) -> "AxonInfo":
+    def _from_parameter_dict_torch(
+        cls, parameter_dict: "torch.nn.ParameterDict"
+    ) -> "AxonInfo":
+        """Returns an axon_info object from a torch parameter_dict."""
+        return cls(**dict(parameter_dict))
+
+    @classmethod
+    def _from_parameter_dict_numpy(cls, parameter_dict: dict[str, Any]) -> "AxonInfo":
         r"""Returns an axon_info object from a parameter_dict."""
         return cls(**parameter_dict)
 
+    @classmethod
+    def from_parameter_dict(
+        cls, parameter_dict: Union[dict[str, Any], "torch.nn.ParameterDict"]
+    ) -> "AxonInfo":
+        if os.environ.get("USE_TORCH"):
+            return cls._from_parameter_dict_torch(parameter_dict)
+        else:
+            return cls._from_parameter_dict_numpy(parameter_dict)
+
 
 class ChainDataType(Enum):
     NeuronInfo = 1
@@ -980,15 +1008,41 @@ def fix_decoded_values(cls, decoded: Dict) -> "SubnetInfo":
             owner_ss58=ss58_encode(decoded["owner"], bittensor.__ss58_format__),
         )
 
-    def to_parameter_dict(self) -> dict[str, Any]:
-        r"""Returns a dict of the subnet info."""
+    def _to_parameter_dict_torch(self) -> "torch.nn.ParameterDict":
+        """Returns a torch tensor of the subnet info."""
+        return maybe_get_torch().nn.ParameterDict(self.__dict__)
+
+    def _to_parameter_dict_numpy(self) -> dict[str, Any]:
+        """Returns a dict of the subnet info."""
         return self.__dict__
 
+    def to_parameter_dict(self) -> Union[dict[str, Any], "torch.nn.ParameterDict"]:
+        if os.environ.get("USE_TORCH"):
+            return self._to_parameter_dict_torch()
+        else:
+            return self._to_parameter_dict_numpy()
+
+    @classmethod
+    def _from_parameter_dict_torch(
+        cls, parameter_dict: "torch.nn.ParameterDict"
+    ) -> "SubnetInfo":
+        """Returns a SubnetInfo object from a torch parameter_dict."""
+        return cls(**dict(parameter_dict))
+
     @classmethod
-    def from_parameter_dict(cls, parameter_dict: dict[str, Any]) -> "SubnetInfo":
+    def _from_parameter_dict_numpy(cls, parameter_dict: dict[str, Any]) -> "SubnetInfo":
         r"""Returns a SubnetInfo object from a parameter_dict."""
         return cls(**parameter_dict)
 
+    @classmethod
+    def from_parameter_dict(
+        cls, parameter_dict: Union[dict[str, Any], "torch.nn.ParameterDict"]
+    ) -> "SubnetInfo":
+        if os.environ.get("USE_TORCH"):
+            return cls._from_parameter_dict_torch(parameter_dict)
+        else:
+            return cls._from_parameter_dict_numpy(parameter_dict)
+
 
 @dataclass
 class SubnetHyperparameters:
@@ -1074,15 +1128,45 @@ def fix_decoded_values(cls, decoded: Dict) -> "SubnetHyperparameters":
             difficulty=decoded["difficulty"],
         )
 
-    def to_parameter_dict(self) -> dict[str, Union[int, float, bool]]:
-        r"""Returns a dict of the subnet hyperparameters."""
+    def _to_parameter_dict_torch(self) -> "torch.nn.ParameterDict":
+        """Returns a torch tensor of the subnet hyperparameters."""
+        return maybe_get_torch().nn.ParameterDict(self.__dict__)
+
+    def _to_parameter_dict_numpy(self) -> dict[str, Union[int, float, bool]]:
+        """Returns a dict of the subnet hyperparameters."""
         return self.__dict__
 
+    def to_parameter_dict(
+        self,
+    ) -> Union[dict[str, Union[int, float, bool]], "torch.nn.ParameterDict"]:
+        if os.environ.get("USE_TORCH"):
+            return self._to_parameter_dict_torch()
+        else:
+            return self._to_parameter_dict_numpy()
+
+    @classmethod
+    def _from_parameter_dict_torch(
+        cls, parameter_dict: "torch.nn.ParameterDict"
+    ) -> "SubnetHyperparameters":
+        """Returns a SubnetHyperparameters object from a torch parameter_dict."""
+        return cls(**dict(parameter_dict))
+
     @classmethod
-    def from_parameter_dict(cls, parameter_dict: dict[str, Any]) -> "SubnetInfo":
-        r"""Returns a SubnetHyperparameters object from a parameter_dict."""
+    def _from_parameter_dict_numpy(
+        cls, parameter_dict: dict[str, Any]
+    ) -> "SubnetHyperparameters":
+        """Returns a SubnetHyperparameters object from a parameter_dict."""
         return cls(**parameter_dict)
 
+    @classmethod
+    def from_parameter_dict(
+        cls, parameter_dict: Union[dict[str, Any], "torch.nn.ParameterDict"]
+    ) -> "SubnetHyperparameters":
+        if os.environ.get("USE_TORCH"):
+            return cls._from_parameter_dict_torch(parameter_dict)
+        else:
+            return cls._from_parameter_dict_numpy(parameter_dict)
+
 
 @dataclass
 class IPInfo:
@@ -1137,15 +1221,43 @@ def fix_decoded_values(cls, decoded: Dict) -> "IPInfo":
             protocol=decoded["ip_type_and_protocol"] & 0xF,
         )
 
-    def to_parameter_dict(self) -> dict[str, Union[str, int]]:
-        r"""Returns a dict of the subnet ip info."""
+    def _to_parameter_dict_torch(self) -> "torch.nn.ParameterDict":
+        """Returns a torch tensor of the subnet info."""
+        return maybe_get_torch().nn.ParameterDict(self.__dict__)
+
+    def _to_parameter_dict_numpy(self) -> dict[str, Union[str, int]]:
+        """Returns a dict of the subnet ip info."""
         return self.__dict__
 
+    def to_parameter_dict(
+        self,
+    ) -> Union[dict[str, Union[str, int]], "torch.nn.ParameterDict"]:
+        if os.environ.get("USE_TORCH"):
+            return self._to_parameter_dict_torch()
+        else:
+            return self._to_parameter_dict_numpy()
+
+    @classmethod
+    def _from_parameter_dict_torch(
+        cls, parameter_dict: "torch.nn.ParameterDict"
+    ) -> "IPInfo":
+        """Returns a IPInfo object from a torch parameter_dict."""
+        return cls(**dict(parameter_dict))
+
     @classmethod
-    def from_parameter_dict(cls, parameter_dict: dict[str, Any]) -> "IPInfo":
-        r"""Returns a IPInfo object from a parameter_dict."""
+    def _from_parameter_dict_numpy(cls, parameter_dict: dict[str, Any]) -> "IPInfo":
+        """Returns a IPInfo object from a parameter_dict."""
         return cls(**parameter_dict)
 
+    @classmethod
+    def from_parameter_dict(
+        cls, parameter_dict: Union[dict[str, Any], "torch.nn.ParameterDict"]
+    ) -> "IPInfo":
+        if os.environ.get("USE_TORCH"):
+            return cls._from_parameter_dict_torch(parameter_dict)
+        else:
+            return cls._from_parameter_dict_numpy(parameter_dict)
+
 
 # Senate / Proposal data
 

bittensor/commands/root.py

@@ -15,16 +15,16 @@
 # OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
 # DEALINGS IN THE SOFTWARE.
 
+import os
 import re
-import numpy as np
 import typing
 import argparse
 import numpy as np
 import bittensor
 from typing import List, Optional, Dict
 from rich.prompt import Prompt
 from rich.table import Table
-from .utils import get_delegates_details, DelegatesDetails
+from .utils import get_delegates_details, DelegatesDetails, maybe_get_torch
 
 from . import defaults
 
@@ -301,7 +301,11 @@ def _run(cli: "bittensor.cli", subtensor: "bittensor.subtensor"):
             f"Boosting weight for netuid {cli.config.netuid} from {prev_weight} -> {new_weight}"
         )
         my_weights[cli.config.netuid] = new_weight
-        all_netuids = np.arange(len(my_weights))
+        all_netuids = (
+            maybe_get_torch().tensor(list(range(len(my_weights))))
+            if os.environ.get("USE_TORCH")
+            else np.arange(len(my_weights))
+        )
 
         bittensor.__console__.print("Setting root weights...")
         subtensor.root_set_weights(
@@ -419,7 +423,11 @@ def _run(cli: "bittensor.cli", subtensor: "bittensor.subtensor"):
         my_weights = root.weights[my_uid]
         my_weights[cli.config.netuid] -= cli.config.amount
         my_weights[my_weights < 0] = 0  # Ensure weights don't go negative
-        all_netuids = np.arange(len(my_weights))
+        all_netuids = (
+            maybe_get_torch().tensor(list(range(len(my_weights))))
+            if os.environ.get("USE_TORCH")
+            else np.arange(len(my_weights))
+        )
 
         subtensor.root_set_weights(
             wallet=wallet,
@@ -520,12 +528,21 @@ def _run(cli: "bittensor.cli", subtensor: "bittensor.subtensor"):
             cli.config.weights = Prompt.ask(f"Enter weights (e.g. {example})")
 
         # Parse from string
-        netuids = np.array(
-            list(map(int, re.split(r"[ ,]+", cli.config.netuids))), dtype=np.int64
+        matched_netuids = list(map(int, re.split(r"[ ,]+", cli.config.netuids)))
+        netuids = (
+            maybe_get_torch().tensor(matched_netuids, dtype=maybe_get_torch().long)
+            if os.environ.get("USE_TORCH")
+            else np.array(matched_netuids, dtype=np.int64)
         )
-        weights = np.array(
-            list(map(float, re.split(r"[ ,]+", cli.config.weights))),
-            dtype=np.float32,
+
+        matched_weights = list(map(float, re.split(r"[ ,]+", cli.config.weights)))
+        weights = (
+            maybe_get_torch().tensor(matched_weights, dtype=maybe_get_torch().float32)
+            if os.environ.get("USE_TORCH")
+            else np.array(
+                matched_weights,
+                dtype=np.float32,
+            )
         )
 
         # Run the set weights operation.

bittensor/dendrite.py

@@ -20,14 +20,16 @@
 from __future__ import annotations
 
 import asyncio
+import os
 import uuid
 import time
 import aiohttp
 import bittensor
 from typing import Union, Optional, List, Union, AsyncGenerator, Any
+from .utils import maybe_get_torch
 
 
-class dendrite:
+class DendriteMixin:
     """
     The Dendrite class represents the abstracted implementation of a network client module.
 
@@ -121,9 +123,6 @@ def __init__(
 
         self._session: Optional[aiohttp.ClientSession] = None
 
-    async def __call__(self, *args, **kwargs):
-        return await self.forward(*args, **kwargs)
-
     @property
     async def session(self) -> aiohttp.ClientSession:
         """
@@ -808,3 +807,18 @@ def __del__(self):
             del dendrite  # This will implicitly invoke the __del__ method and close the session.
         """
         self.close_session()
+
+
+if os.environ.get("USE_TORCH"):
+    class dendrite(maybe_get_torch().nn.module, DendriteMixin):
+        def __init__(self):
+            maybe_get_torch().nn.module.__init__(self)
+            DendriteMixin.__init__(self)
+else:
+    class dendrite(DendriteMixin):
+        def __init__(self):
+            DendriteMixin.__init__(self)
+
+        async def __call__(self, *args, **kwargs):
+            return await self.forward(*args, **kwargs)
+