Skip to content

Commit

Permalink
Adds Device.getDevices() for all Device (#2820)
Browse files Browse the repository at this point in the history
This improves upon the creation of MultiDevice in #2819 by moving the getDevices
function to the main Device class. This can simplify the usage of something
which is potentially a MultiDevice and make it easier to check for the presence
of a MultiDevice.
  • Loading branch information
zachgk authored and frankfliu committed Apr 26, 2024
1 parent dda5f75 commit 36f4e92
Show file tree
Hide file tree
Showing 2 changed files with 13 additions and 5 deletions.
17 changes: 12 additions & 5 deletions api/src/main/java/ai/djl/Device.java
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
import ai.djl.engine.Engine;

import java.util.Arrays;
import java.util.Collections;
import java.util.Comparator;
import java.util.List;
import java.util.Map;
Expand Down Expand Up @@ -162,6 +163,15 @@ public boolean isGpu() {
return Type.GPU.equals(deviceType);
}

/**
* Returns the sub devices if present (such as a {@link MultiDevice}), otherwise this.
*
* @return the sub devices if present (such as a {@link MultiDevice}), otherwise this.
*/
public List<Device> getDevices() {
return Collections.singletonList(this);
}

/** {@inheritDoc} */
@Override
public String toString() {
Expand Down Expand Up @@ -276,11 +286,8 @@ public MultiDevice(List<Device> devices) {
this.devices = devices;
}

/**
* Returns the sub devices.
*
* @return the sub devices
*/
/** {@inheritDoc} */
@Override
public List<Device> getDevices() {
return devices;
}
Expand Down
1 change: 1 addition & 0 deletions api/src/test/java/ai/djl/DeviceTest.java
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,7 @@ public void testDevice() {
System.setProperty("test_key", "test");
Engine.debugEnvironment();

Assert.assertEquals(1, Device.cpu().getDevices().size());
Assert.assertEquals(2, new MultiDevice(Device.gpu(1), Device.gpu(2)).getDevices().size());
}

Expand Down

0 comments on commit 36f4e92

Please sign in to comment.