Lines Matching refs:device
15 """Tests for device function for replicated training."""
37 with ops.device(
39 with ops.device("/cpu:0"):
42 with ops.device("/cpu:0"):
44 self.assertDeviceEqual("/job:ps/task:0/cpu:0", v.device)
45 self.assertDeviceEqual("/job:ps/task:0/cpu:0", v.initializer.device)
46 self.assertDeviceEqual("/job:ps/task:1", w.device)
47 self.assertDeviceEqual("/job:ps/task:1", w.initializer.device)
48 self.assertDeviceEqual("/job:worker/cpu:0", a.device)
51 with ops.device(
54 self.assertDeviceEqual("/job:ps/task:0", v.device)
57 with ops.device(
62 self.assertDeviceEqual("/job:ps/task:0", v.device)
63 self.assertDeviceEqual("/job:ps/task:0", v.initializer.device)
64 self.assertDeviceEqual("/job:ps/task:1", w.device)
65 self.assertDeviceEqual("/job:ps/task:1", w.initializer.device)
66 self.assertDeviceEqual("/job:worker", a.device)
69 with ops.device(
72 with ops.device("/job:moon"):
74 with ops.device("/job:ps"): # Explicit PS job will get task set.
77 self.assertDeviceEqual("/job:ps/task:0", v.device)
78 self.assertDeviceEqual("/job:ps/task:0", v.initializer.device)
79 self.assertDeviceEqual("/job:moon", w.device)
80 self.assertDeviceEqual("/job:moon", w.initializer.device)
81 self.assertDeviceEqual("/job:ps/task:1", x.device)
82 self.assertDeviceEqual("/job:ps/task:1", x.initializer.device)
83 self.assertDeviceEqual("/job:worker", a.device)
86 with ops.device(
89 with ops.device("/job:moon"):
92 self.assertDeviceEqual("/cpu:0", v.device)
93 self.assertDeviceEqual("/cpu:0", v.initializer.device)
94 self.assertDeviceEqual("/job:moon/cpu:0", w.device)
95 self.assertDeviceEqual("/job:moon/cpu:0", w.initializer.device)
96 self.assertDeviceEqual("/job:worker", a.device)
99 with ops.device(
103 with ops.device("/job:ps"): # Won't assign task when merge_devices=False.
106 self.assertDeviceEqual("/job:ps/task:0", v.device)
107 self.assertDeviceEqual("/job:ps/task:0", v.initializer.device)
108 self.assertDeviceEqual("/job:ps", w.device)
109 self.assertDeviceEqual("/job:ps", w.initializer.device)
110 self.assertDeviceEqual("/job:worker", a.device)
113 with ops.device(
119 self.assertDeviceEqual("/job:ps/task:0", v.device)
120 self.assertDeviceEqual("/job:ps/task:0", v.initializer.device)
121 self.assertDeviceEqual("/job:ps/task:1", w.device)
122 self.assertDeviceEqual("/job:ps/task:1", w.initializer.device)
123 self.assertDeviceEqual("/job:worker", a.device)
126 with ops.device(
132 self.assertDeviceEqual("/job:ps/task:0", v.device)
133 self.assertDeviceEqual("/job:ps/task:0", v.initializer.device)
134 self.assertDeviceEqual("/job:ps/task:1", w.device)
135 self.assertDeviceEqual("/job:ps/task:1", w.initializer.device)
136 self.assertDeviceEqual("/job:worker", a.device)
144 with ops.device(
152 self.assertDeviceEqual("/job:moon/task:0", v.device)
153 self.assertDeviceEqual("/job:moon/task:0", v.initializer.device)
154 self.assertDeviceEqual("/job:moon/task:1", w.device)
155 self.assertDeviceEqual("/job:moon/task:1", w.initializer.device)
156 self.assertDeviceEqual("/job:sun", a.device)
164 with ops.device(
172 self.assertDeviceEqual("/job:moon/task:0/cpu:0", v.device)
173 self.assertDeviceEqual("/job:moon/task:0/cpu:0", v.initializer.device)
174 self.assertDeviceEqual("/job:moon/task:1/cpu:0", w.device)
175 self.assertDeviceEqual("/job:moon/task:1/cpu:0", w.initializer.device)
176 self.assertDeviceEqual("/job:sun", a.device)
