Op เพื่อแลกเปลี่ยนข้อมูลระหว่างแบบจำลอง TPU
ในแต่ละเรพลิกา อินพุตจะถูกแบ่งออกเป็นบล็อก `split_count` ตาม 'split_daiment` และส่งไปยังเรพลิกาอื่นๆ ที่กำหนด group_assignment หลังจากได้รับ `split_count` - 1 บล็อกจากแบบจำลองอื่น เราจะต่อบล็อกต่างๆ ตาม `concat_ dimensions` เป็นเอาต์พุต
ตัวอย่างเช่น สมมติว่ามีแบบจำลอง TPU 2 ตัว: แบบจำลอง 0 ได้รับอินพุต: `[[A, B]]` แบบจำลอง 1 ได้รับอินพุต: `[[C, D]]`
group_signment=`[[0, 1]]` concat_size=0 split_ dimensions=1 split_count=2
เอาต์พุตของแบบจำลอง 0: `[[A], [C]]` เอาต์พุตของแบบจำลอง 1: `[[B], [D]]`
วิธีการสาธารณะ
| เอาท์พุต <T> | เป็นเอาท์พุต () ส่งกลับค่าแฮนเดิลสัญลักษณ์ของเทนเซอร์ |
| คงที่ <T> AllToAll <T> | |
| เอาท์พุต <T> | เอาท์พุท () ผลการแลกเปลี่ยน. |
วิธีการสืบทอด
วิธีการสาธารณะ
เอาท์ พุท สาธารณะ <T> asOutput ()
ส่งกลับค่าแฮนเดิลสัญลักษณ์ของเทนเซอร์
อินพุตสำหรับการดำเนินการ TensorFlow คือเอาต์พุตของการดำเนินการ TensorFlow อื่น วิธีการนี้ใช้เพื่อรับหมายเลขอ้างอิงสัญลักษณ์ที่แสดงถึงการคำนวณอินพุต
สร้าง AllToAll <T> แบบคงที่สาธารณะ (ขอบเขต ขอบเขต , ตัวดำเนินการ <T> อินพุต, ตัวดำเนินการ <Integer> groupAssignment, Long concatDimension, Long splitDimension, Long splitCount)
วิธีการจากโรงงานเพื่อสร้างคลาสที่รวมการดำเนินการ AllToAll ใหม่
พารามิเตอร์
| ขอบเขต | ขอบเขตปัจจุบัน |
|---|---|
| ป้อนข้อมูล | อินพุตท้องถิ่นเพื่อผลรวม |
| การมอบหมายกลุ่ม | เทนเซอร์ int32 ที่มีรูปร่าง [num_groups, num_replicas_per_group] `group_assignment[i]` แสดงถึงรหัสจำลองในกลุ่มย่อยที่ ith |
| เชื่อมต่อมิติ | หมายเลขมิติที่จะต่อกัน |
| แยกมิติ | หมายเลขมิติที่จะแยก |
| แยกนับ | จำนวนการแยก จำนวนนี้ต้องเท่ากับขนาดกลุ่มย่อย(group_assignment.get_shape()[1]) |
การส่งคืน
- อินสแตนซ์ใหม่ของ AllToAll