zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-09-02 13:54:44 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Typing the dropout. (#1479)

Browse Source
This commit is contained in:
LastWhisper
2024-10-15 21:45:46 +08:00
committed by GitHub
parent 0ab8e099e8
commit 2b8ace6a03
1 changed files with 6 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
python/mlx/nn/layers/dropout.py
View File

@@ -23,10 +23,10 @@ class Dropout(Module):
self._p_1 = 1 - p
def _extra_repr(self):
def _extra_repr(self) -> str:
return f"p={1-self._p_1}"
def __call__(self, x):
def __call__(self, x: mx.array) -> mx.array:
if self._p_1 == 1 or not self.training:
return x
@@ -66,10 +66,10 @@ class Dropout2d(Module):
self._p_1 = 1 - p
def _extra_repr(self):
def _extra_repr(self) -> str:
return f"p={1-self._p_1}"
def __call__(self, x):
def __call__(self, x: mx.array) -> mx.array:
if x.ndim not in (3, 4):
raise ValueError(
f"Received input with {x.ndim} dimensions. Expected 3 or 4 dimensions."
@@ -115,10 +115,10 @@ class Dropout3d(Module):
self._p_1 = 1 - p
def _extra_repr(self):
def _extra_repr(self) -> str:
return f"p={1-self._p_1}"
def __call__(self, x):
def __call__(self, x: mx.array) -> mx.array:
if x.ndim not in (4, 5):
raise ValueError(
f"Received input with {x.ndim} dimensions. Expected 4 or 5 dimensions."